#ifndef MYAMG_BICGSTABSOLVER_CUH
#define MYAMG_BICGSTABSOLVER_CUH
#include <Solver/Solver.cuh>
#include <utils/cudaProfiler.cuh>

namespace myamg {
    template<typename Integer, typename Real, typename Allocator>
    class BiCGStabSolver : public Solver<Integer, Real, Allocator, BiCGStabSolver<Integer, Real, Allocator> > {
    public:
        using matrixType = typename Solver<Integer, Real, Allocator, BiCGStabSolver>::matrixType;
        using vectorType = typename Solver<Integer, Real, Allocator, BiCGStabSolver>::vectorType;
        using integerType = typename Solver<Integer, Real, Allocator, BiCGStabSolver>::integerType;
        using realType = typename Solver<Integer, Real, Allocator, BiCGStabSolver>::realType;
        using allocatorType = typename Solver<Integer, Real, Allocator, BiCGStabSolver>::allocatorType;

        explicit BiCGStabSolver(const int maxIters = 500, const Real tol = 1e-8, const bool print = false)
            : maxIterations(maxIters), tolerance(tol), printStats(print) {
        }

        template<typename Precond>
        void applyImpl(const matrixType &A, Precond &P, const vectorType &b, vectorType &x, int &iter,
                       double &relres) {
            auto n = A.numRows();
            dim3 block(512);
            dim3 grid((n + block.x - 1) / block.x);
            vectorType r(n);
            vectorType p(n);
            vectorType q(n);
            vectorType v(n);
            vectorType s(n);
            vectorType t(n);

            Real normb, res, alpha, beta, omega, rho, rho1;

            if (printStats) {
                std::cout << "Iter\tResidual\n";
                std::cout << "--------------------------------------------------\n";
            }

            iter = 0;

            // r = b - Ax
            if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                calcResidualHost(r.values(), A.rowPtr(), A.colIndex(), A.values(), b.values(), x.values(), n);
            } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                calcResidualDevice<<<grid, block>>>(r.values(), A.rowPtr(), A.colIndex(), A.values(), b.values(),
                                                    x.values(), n);
            } else {
                std::cerr << "Unsupported allocator type in BiCGStabSolver." << std::endl;
            }

            vectorType r0(r);
            rho = vectorDotProduct(r, r);
            res = std::sqrt(rho);
            normb = std::sqrt(vectorDotProduct(b, b));

            while (true) {
                if (printStats) {
                    std::cout << iter << "\t" << res << "\n";
                }

                if (res / normb <= tolerance || iter == maxIterations) {
                    break;
                }

                if (iter == 0) {
                    p = r;
                } else {
                    beta = (alpha / omega) * (rho / rho1);
                    // p = r + beta * (p - omega * v)
                    if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                        updatePHost(p.values(), r.values(), v.values(), omega, beta, n);
                    } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                        updatePDevice<<<grid, block>>>(p.values(), r.values(), v.values(), omega, beta, n);
                    } else {
                        std::cerr << "Unsupported allocator type in BiCGStabSolver." << std::endl;
                    }
                }

                // Aq = p，使用预条件子求解q
                P.apply(p, q);

                // v = Aq
                matMulVec(A, q, v);

                // alpha = rho / (r0, v)
                alpha = rho / vectorDotProduct(r0, v);

                // r = r - alpha * v
                if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                    BiCGStabUpdateRHost(r.values(), v.values(), alpha, n);
                } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                    BiCGStabUpdateRDevice<<<grid, block>>>(r.values(), v.values(), alpha, n);
                } else {
                    std::cerr << "Unsupported allocator type in BiCGStabSolver." << std::endl;
                }
                P.apply(r, s);

                // t = As
                matMulVec(A, s, t);

                // omega = (t, r) / (t, t)
                omega = vectorDotProduct(t, r) / vectorDotProduct(t, t);

                // x = x + alpha * q + omega * s
                if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                    BiCGStabUpdateXHost(x.values(), q.values(), s.values(), alpha, omega, n);
                } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                    BiCGStabUpdateXDevice<<<grid, block>>>(x.values(), q.values(), s.values(), alpha, omega, n);
                } else {
                    std::cerr << "Unsupported allocator type in BiCGStabSolver." << std::endl;
                }

                // r = r - omega * t
                if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                    BiCGStabUpdateRHost(r.values(), t.values(), omega, n);
                } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                    BiCGStabUpdateRDevice<<<grid, block>>>(r.values(), t.values(), omega, n);
                } else {
                    std::cerr << "Unsupported allocator type in BiCGStabSolver." << std::endl;
                }

                rho1 = rho;
                rho = vectorDotProduct(r0, r);
                res = std::sqrt(vectorDotProduct(r, r));

                ++iter;
            }

            relres = res / normb;
            if (printStats) {
                std::cout << "--------------------------------------------------\n";
                std::cout << "Iterations: " << iter << "\n";
                std::cout << "Final Relative Residual: " << relres << "\n";
            }
        }

        int maxIterations; //!< 最大迭代数，默认为500
        Real tolerance; //!< 停止误差，默认为1e-8
        bool printStats; //!< 是否打印迭代信息，默认为false
    };
}

#endif //MYAMG_BICGSTABSOLVER_CUH
